import sys, os, time, getopt, subprocess, tempfile
import os
import sys
sys.path.append(os.path.dirname(__file__))  # 加入当前目录（parsec/）

from parsec_platform import *

abspath = lambda d: os.path.abspath(os.path.join(d))
HOME = abspath(os.path.dirname(__file__))

__allbenchmarks = None

def allbenchmarks():
    global __allbenchmarks
    if not __allbenchmarks:
        try:
            benchmarks = subprocess.Popen(
                [f'{HOME}/parsec-2.1/bin/parsecmgmt', '-a', 'info'],
                stdout=subprocess.PIPE
            ).communicate()
            lines = benchmarks[0].decode().split('\n')
            benchmarks = [line[15:].split(' ')[0] for line in lines if line.startswith('[PARSEC]     - ') and (line.endswith(' (apps)') or line.endswith(' (kernels)'))]
            __allbenchmarks = sorted(benchmarks)
        except OSError:
            return None
    return __allbenchmarks

def allinputs():
    return [f[:-8] for f in os.listdir(f'{HOME}/parsec-2.1/config') if f.endswith('.runconf')]

def log2(n):
    log2n = -1
    while n:
        n >>= 1
        log2n += 1
    return log2n

class Program:

    def __init__(self, program, nthreads, inputsize, benchmark_options=[]):
        if program not in allbenchmarks():
            raise ValueError(f"Invalid benchmark {program}")
        if inputsize not in allinputs():
            if inputsize in ('small', 'medium', 'large'):
                inputsize = 'sim' + inputsize
            else:
                raise ValueError(f"Invalid input size {inputsize}")

        self.program = program
        self.nthreads = int(nthreads)
        self.nthreads_force = 'force_nthreads' in benchmark_options
        self.inputsize = inputsize
        self.openmp = program in ('freqmine',)

        for option in benchmark_options:
            if option.startswith('pinthreads'):
                if '=' in option:
                    corelist = option.split('=')[1]
                else:
                    corelist = ','.join(map(str, range(nthreads)))
                os.environ['PARMACS_PINTHREADS'] = corelist

        nthreads = self.get_nthreads()

        if self.program == 'facesim':
            nthreads_supported = (1, 2, 3, 4, 6, 8, 16, 32, 64, 128)
            if nthreads not in nthreads_supported:
                raise ValueError(f"Benchmark {self.program} only supports running with {nthreads_supported} threads")
        elif self.program == 'fluidanimate':
            if nthreads != 1 << log2(nthreads):
                raise ValueError(f"Benchmark {self.program}: number of threads must be power of two")

    def get_nthreads(self):
        if self.nthreads_force:
            return self.nthreads

        if self.program == 'blackscholes':
            nthreads = self.nthreads - 1
        elif self.program == 'bodytrack':
            nthreads = self.nthreads - 2
        elif self.program == 'facesim':
            nthreads = self.nthreads
        elif self.program == 'ferret':
            nthreads = (self.nthreads - 2) // 4
        elif self.program == 'fluidanimate':
            nthreads = 1 << log2(self.nthreads - 1)
        elif self.program == 'swaptions':
            nthreads = self.nthreads - 1
        elif self.program == 'canneal':
            nthreads = self.nthreads - 1
        elif self.program == 'raytrace':
            nthreads = self.nthreads - 1
        elif self.program == 'dedup':
            nthreads = self.nthreads // 4
        elif self.program == 'streamcluster':
            nthreads = self.nthreads - 1
        elif self.program == 'vips':
            nthreads = self.nthreads - 2
        else:
            nthreads = self.nthreads

        if nthreads < 1:
            raise ValueError(f"Benchmark {self.program} needs more cores")

        return nthreads

    def run(self, graphitecmd, postcmd=''):
        if postcmd != '':
            sys.stderr.write('PARSEC Error: postcmd not supported\n')
            return 1

        flags = []
        rundir = tempfile.mkdtemp()

        if self.program == 'facesim':
            inputfile = f'{HOME}/parsec-2.1/pkgs/apps/facesim/inputs/input_{self.inputsize}.tar'
            if not os.path.exists(inputfile):
                print(f'Cannot find input file {inputfile}')
                sys.exit(-1)
            flags.append('-k')
            os.system(f'rm -r {rundir}/apps/facesim/run')
            os.system(f'mkdir -p {rundir}/apps/facesim/run/Storytelling/output')
            os.system(f'tar xvf {inputfile} -C {rundir}/apps/facesim/run')

        if self.openmp:
            os.environ['OMP_NUM_THREADS'] = str(self.get_nthreads())

        proc = subprocess.Popen([
            f'{HOME}/parsec-2.1/bin/parsecmgmt',
            '-a', 'run',
            '-p', self.program,
            '-c', PLATFORM,
            '-i', self.inputsize,
            '-n', str(self.get_nthreads()),
            '-s', graphitecmd,
            '-d', rundir
        ] + flags)
        proc.communicate()

        os.system(f'rm -r {rundir}')
        return proc.returncode

    def rungraphiteoptions(self):
        return ''
